import os,re
import cv2
import easyocr
from tqdm import tqdm
import pandas as pd
import numpy as np
from multiprocessing import Pool, cpu_count

def read_image_safe(path):
    """支持中文路径的图片读取"""
    try:
        with open(path, 'rb') as f:
            img_bytes = f.read()
        img_buffer = np.frombuffer(img_bytes, dtype=np.uint8)
        img = cv2.imdecode(img_buffer, cv2.IMREAD_COLOR)
        return img
    except:
        return None

def process_single_image(image_path):
    """单个图片处理函数（每个子进程独立执行）"""
    try:
        # 每个进程独立初始化EasyOCR
        reader = easyocr.Reader(['ch_sim', 'en'], gpu=False)
        
        # 读取图片
        img = read_image_safe(image_path)
        if img is None:
            return {
                'image_path': image_path,
                'text': '[ERROR] 图片读取失败'
            }
        
        # 转换为RGB并执行OCR
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        result = reader.readtext(img, detail=0)
        return {
            'image_path': image_path,
            'text': ' '.join(result)
        }
    except Exception as e:
        return {
            'image_path': image_path,
            'text': f'[ERROR] {str(e)}'
        }

def batch_ocr(image_folder, output_file='output.csv'):
    # 收集所有图片路径
    image_list = [
        os.path.join(image_folder, f)
        for f in os.listdir(image_folder)
        if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))
    ]

    # 多进程处理
    with Pool(cpu_count()) as pool:
        results = list(tqdm(
            pool.imap(process_single_image, image_list),
            total=len(image_list),
            desc="多进程OCR处理"
        ))
    
    # 保存结果到CSV
    df = pd.DataFrame(results)
    df.to_csv(output_file, index=False, encoding='utf-8-sig')
    print(f"✅ 结果已保存至 {output_file}")

if __name__ == "__main__":
    # 注意：路径使用原始字符串避免转义问题
    batch_ocr(image_folder=r'D:/Desk/0312-炳哥尺寸转换/尺寸标注/标注信息')